df <- readRDS('Embed/simulatedCrossSectionalDyadicData.rds')
print_df(head(df))

df_apim <- df %>%
  # Reshaping to Actor-Partner format (4-field)
  reshape_dyadic_data(
    person_id = 'userID',
    dyad_id = 'coupleID',
    vars_to_reshape = 'communication'
  )  %>%
  mutate(
    # Optional: grand-mean centering
    communication_actor_gmc = communication_actor - mean(communication_actor),
    communication_partner_gmc = communication_partner - mean(communication_partner),
    
    # Create Dummy-Variables for male and female
    is_male = ifelse(gender == 1, 1, 0),
    is_female = ifelse(gender == 2, 1, 0),
  ) 

print_df(head(df_apim))

formula <- bf(
  satisfaction ~ 
    
    # Remove global intercept, introduce male and female intercepts.
    0 + is_male + is_female +
    
    # Actor effect for male and female
    communication_actor_gmc:is_male + 
    communication_actor_gmc:is_female + 
    
    # Partner effect for male and female
    communication_partner_gmc:is_male + 
    communication_partner_gmc:is_female +
    
    # Option 1: Random intercept for male and female (correlated), not identifiable in the cross-sectional case:
    # (0 + is_male + is_female | coupleID)
    
    # Option 2: Too restrictive, basically modelling Exchangeable dyads with perfect correlation
    # (1 | coupleID) 
    
    # Solution: use compound symmetry of the residuals
    cosy(gr = coupleID),
  
  # Two distinct residual variances for males vs females:
  sigma = ~ 0 + is_male + is_female
)

priors <- c(
  # prior(normal(2, 3), class = "Intercept"),
  prior(normal(0, 5), class = "b"),
  prior(student_t(3, 0, 1.5), class = "b", dpar = "sigma"),
  prior(beta(1, 3), class = "cosy")
)

model_dist_apim <- brm(
  formula = formula, 
  data = df_apim,
  family = gaussian(link = identity),
  prior = priors,
  chains = 4,
  cores = 4,
  iter = 2000,
  warmup = 1000, 
  seed = 123,
  file = file.path('brms_cache', 'example1_dist_apim') # Cache the model
)

rstan::check_hmc_diagnostics(model_dist_apim$fit)

loo::pareto_k_table(loo(model_dist_apim))

plot(model_dist_apim, ask = FALSE)

pp_check(model_dist_apim, 'dens_overlay_grouped', group = 'is_male')

pp_check(model_dist_apim, 'ecdf_overlay_grouped', group = 'is_male')

pp_check(model_dist_apim, type = "loo_pit_overlay")

# Custom function to make DHARMa work with brms (see file 'Functions')
DHARMa.check_brms(model_dist_apim)

df_long_apim <- bmlm::isolate(
  df_long_apim, 
  by = 'userID', # NOT coupleID
  value = c('provided_support_actor', 'provided_support_partner'),
  which = 'both' 
) %>%
  rename(
    provided_support_actor_cwp = provided_support_actor_cw,
    provided_support_partner_cwp = provided_support_partner_cw,
    provided_support_actor_cbp = provided_support_actor_cb,
    provided_support_partner_cbp = provided_support_partner_cb,
  )

formula <- bf(
  daily_closeness ~ 
    
    # Intercepts
    0 + is_male + is_female +
    
    # Time Slopes
    is_male:diaryday +
    is_female:diaryday +
    
    # Between-Person APIM
    is_male:provided_support_actor_cbp + 
    is_male:provided_support_partner_cbp +
    
    is_female:provided_support_actor_cbp + 
    is_female:provided_support_partner_cbp +
  
    # Within-Person APIM
    is_male:provided_support_actor_cwp + 
    is_male:provided_support_partner_cwp +
    
    is_female:provided_support_actor_cwp + 
    is_female:provided_support_partner_cwp +
    
    # Accounting for non-independence between partners' means and trajectories 
    # and effect sensitivities via random effects:
    (0 + is_male + is_female +  # Random Intercepts and their correlations
       
       is_male:diaryday + is_female:diaryday +
       
       is_male:provided_support_actor_cwp + 
       is_male:provided_support_partner_cwp +
       is_female:provided_support_actor_cwp + 
       is_female:provided_support_partner_cwp | coupleID ) + 
    
    
    # Accounting for daily non-independence
    
    # Option 1: model correlated residuals on the day with a 
    # compound symmetry structure:
    # cosy(time = diaryday, gr = coupleID),
    
    # Option 2: model daily coupling / common shocks with a random effect:
    (1 | coupleID:diaryday) +
    
    
    # Modelling residuals as autocorrelated within each person (order 1)
    ar(time = diaryday, gr = coupleID:userID, p = 1),
    
  sigma ~ 0 + is_male + is_female   # heteroscedastic residuals
)

priors <- c(
  prior(normal(0, 5), class = "b"),
  prior(student_t(3, 0, 1.5), class = "sd"),
  prior(normal(0.5, 0.3), class = "ar"),
  prior(student_t(3, 0, 1.5), class ="b", dpar = "sigma")
)


model_dist_apim_long <- brm(
  formula = formula, 
  data = df_long_apim,
  family = gaussian(link = identity),
  prior = priors,
  chains = 4,
  cores = 4,
  iter = 2000,
  warmup = 1000, 
  seed = 123,
  file = file.path('brms_cache', 'example1_dist_apim_long') # Cache the model
)

print_df(head(df_apim))

formula <- bf(
  satisfaction ~ 1 + 
    communication_actor_gmc + communication_partner_gmc + 
    
    # Option 1: A single couple level random intercept. Too restrictive.
    # Would impose perfect correlation between partners in each couple
    # (1 | coupleID)
    
    cosy(gr = coupleID)
  
  # Note: no need to model separate sigmas for each partner.
  # Homogeneous residual variance is estimated: 
  # Implied: sigma = ~ 1
)

priors <- c(
  prior(normal(2, 10), class = "Intercept"),
  prior(normal(0, 5), class = "b"),
  prior(student_t(3, 0, 1.5), class = "sigma"),
  prior(beta(1, 3), class = "cosy")
)

model_ind_apim <- brm(
  formula = formula, 
  data = df_apim,
  family = gaussian(link = identity),
  prior = priors,
  chains = 4,
  cores = 4,
  iter = 2000,
  warmup = 1000, 
  seed = 123,
  file = file.path('brms_cache', 'example1_ind_apim') # Cache the model
)

a <- loo_compare(
  loo(model_ind_apim), 
  loo(model_dist_apim)
)
print(a)

report::report(a)

df_dim <- df_apim2 %>%
  group_by(coupleID) %>%
  mutate(
    communication_cm = mean(communication, na.rm = TRUE),
    communication_cwc= communication - communication_cm
  ) %>%
  ungroup() %>%
  mutate(
    communication_cbc= communication_cm - mean(communication_cm, na.rm = TRUE)
  ) %>% 
  select(c('userID', 'coupleID', 'satisfaction', 'communication_cwc', 'communication_cbc'))

print_df(head(df_dim))

formula <- bf(
  satisfaction ~ 1 + 
    communication_cbc+ communication_cwc + 
    cosy(gr = coupleID) 
)

priors <- c(
  prior(normal(2, 10), class = "Intercept"),
  prior(normal(0, 5), class = "b"),
  prior(student_t(3, 0, 1.5), class = "sigma"),
  prior(beta(1, 3), class = "cosy")
)

model_ind_dim <- brm(
  formula = formula, 
  data = df_dim,
  family = gaussian(link = identity),
  prior = priors,
  chains = 4,
  cores = 4,
  iter = 2000,
  warmup = 1000, 
  seed = 123,
  file = file.path('brms_cache', 'example1_ind_dim') # Cache the model
)

a <- hypothesis(
  model_ind_apim, 
  "communication_actor_gmc + communication_partner_gmc = 0"
)

a$hypothesis[,c(2,3,4,5)]

df_long_apim2 <- df_long_apim %>%
  select('coupleID', 'userID', 'diaryday', 'daily_closeness', 
         'provided_support_actor_cwp', 'provided_support_partner_cwp', 
         'provided_support_actor_cbp', 'provided_support_partner_cbp')

print_df(print_couple_preview(df_long_apim2, couple_id = "31"))  

df_long_dim <- df_long_apim2 %>%
  mutate(
    
    # Between Person Level DIM (centering between couples and within couple between person)
    provided_support_cbc = (provided_support_actor_cbp + provided_support_partner_cbp) / 2,
    provided_support_cwcbp = (provided_support_actor_cbp - provided_support_partner_cbp) / 2,
    
    # Within Person Level DIM (couple mean and deviation on each day)
    provided_support_cwp_mean = (provided_support_actor_cwp + provided_support_partner_cwp) / 2,
    provided_support_cwp_halfdiff = (provided_support_actor_cwp - provided_support_partner_cwp) / 2
  ) 

df_long_dim <- df_long_dim %>%
  group_by(coupleID) %>%
  mutate(
    base = ifelse(userID == min(userID), 1, -1),
    flip = 1 - 2 * rbinom(1, 1, 0.5),   # yields +1 or -1 once per couple
    Idiff = base * flip
  ) %>%
  ungroup() %>%
  relocate(Idiff, .after = userID) %>%
  select(-base, -flip)


print_df(print_couple_preview(df_long_dim, couple_id = "31"))  

formula <- bf(
  daily_closeness ~ 1 + 
    
    diaryday +
  
    # Within-person APIM
    provided_support_actor_cwp + 
    provided_support_partner_cwp +
    # Equivalent to within-person DIM:
    # provided_support_cwp_mean + 
    # provided_support_cwp_halfdiff +
  
    # Between-person APIM
    provided_support_actor_cbp +
    provided_support_partner_cbp +
    # Equivalent to between-person DIM
    # provided_support_cwcbp+
    # provided_support_cbc +
    
    # Dyad-Level intercept and slopes for time-varying predictors
    (1 + diaryday + provided_support_actor_cwp + provided_support_partner_cwp | coupleID) +
    # Both partners' deviations from these dyad-level means and slopes
    (0 + Idiff + 
         I(Idiff * diaryday) + 
         I(Idiff * provided_support_actor_cwp) + 
         I(Idiff * provided_support_partner_cwp) | coupleID) +
    
    # Accounting for same-day shocks/coupling
    (1 | coupleID:diaryday) +
    # alternatively and equivalently if no autocorrelation needed:
    # cosy(time = diaryday, gr = coupleID) 
    
    # Autocorrelated residuals
    ar(time = diaryday, gr = coupleID:userID, p = 1)
  
  # Again, no need to model heterogeneous residual variances (sigma)
  # Implied: sigma = ~ 1
)

priors <- c(
  prior(normal(2, 10), class = "Intercept"),
  prior(normal(0, 5), class = "b"),
  prior(student_t(3, 0, 1.5), class = "sd"),
  prior(lkj(2), class = "cor"), # correlation prior for random effect matrix
  prior(normal(0.5, 0.3), class = "ar"),
  prior(student_t(3, 0, 1.5), class = "sigma")
)

model_apim_ind_long <- brm(
  formula = formula, 
  data = df_long_dim,
  family = gaussian(link = identity),
  prior = priors,
  chains = 4,
  cores = 4,
  iter = 2000,
  warmup = 1000, 
  seed = 123,
  file = file.path('brms_cache', 'model_apim_ind_long_both')
)

# Custom function 
rotate_apim_covariance <- function(
    fit,
    gr = "coupleID",
    Idiff = "Idiff"
) {
  random_effects <- fit$ranef
  random_effects <- random_effects[random_effects$group == gr,]

  SUM <- random_effects$coef[!grepl(Idiff, random_effects$coef)]
  DIFF <- random_effects$coef[grepl(Idiff, random_effects$coef)]
  
  varcor <- VarCorr(fit)[[gr]]
  
  within_person_covariance_matrix <- varcor$cov[SUM,'Estimate',SUM] + varcor$cov[DIFF,'Estimate',DIFF]
  cross_person_covariance_matrix <- varcor$cov[SUM,'Estimate',SUM] - varcor$cov[DIFF,'Estimate',DIFF]
  
  p <- nrow(within_person_covariance_matrix)
  
  Full <- rbind(
    cbind(within_person_covariance_matrix,  cross_person_covariance_matrix),
    cbind(cross_person_covariance_matrix, within_person_covariance_matrix)
  )
  
  # nice labels
  base <- SUM
  rn <- c(paste0("PartnerA_", base), paste0("PartnerB_", base))
  colnames(Full) <- rn
  rownames(Full) <- rn
  
  return(list(within_person_covariance_matrix=within_person_covariance_matrix, cross_person_covariance_matrix=cross_person_covariance_matrix, full_covariance_matrix=Full))
}

a <- rotate_apim_covariance(model_apim_ind_long)

print_df(a$within_person_covariance_matrix)

print_df(a$cross_person_covariance_matrix)

print_df(a$full_covariance_matrix)